import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(project_root)
import json
import argparse
import matplotlib
matplotlib.use('Agg')

import numpy as np
import matplotlib.pyplot as plt

from concurrent.futures import ThreadPoolExecutor
from PIL import Image, ImageDraw

from utils import *

X_RANGE = None
Y_RANGE = None

def load_frames_info(file_path):
    mat_data = read_mat_file(file_path)
    if mat_data:
        start_frame_llh = mat_data.get('StartFrameLLH', None)
        start_frame_prh = mat_data.get('StartFrameRPY', None)

        lat = start_frame_llh[0][0]
        lon = start_frame_llh[0][1]
        h = start_frame_llh[0][2] / 1000
        pitch = start_frame_prh[0][0]
        roll = start_frame_prh[0][1]
        heading = start_frame_prh[0][2]
        frames_info = {
            "filename": os.path.splitext(os.path.basename(file_path))[0],
            "lat": lat,
            "lon": lon,
            "h": h,
            "pitch": pitch,
            "roll": roll,
            "heading": heading
        }
    return frames_info

def plot_polygons(X, Y):
    plt.plot(X, Y, marker='o')
    plt.legend()
    plt.grid(True)
    plt.show()
    
def create_mask(frame, lat0, lon0, h0, polygons, x_range, y_range, resolution_multiplier=10):
    x_min, x_max = x_range
    y_min, y_max = y_range

    lat, lon, h = frame["lat"], frame["lon"], frame["h"]
    heading = frame["heading"]

    mask_width = int((x_max - x_min) * resolution_multiplier)
    mask_height = int((y_max - y_min) * resolution_multiplier)
    mask_image = Image.new("L", (mask_width, mask_height), 255)
    draw = ImageDraw.Draw(mask_image)

    for polygon in polygons:
        polygon = np.array(polygon)
        
        X, Y, _ = global_to_local_coords(
            polygon[:, 0], polygon[:, 1], np.zeros(polygon.shape[0]),
            lat, lon, h, 0, 0, heading, lat0, lon0, h0
        )

        # Splot_polygons(X, Y)

        polygon_coords = []
        for x, y in zip(X, Y):
            pixel_x = int(round((x - x_min) * resolution_multiplier))
            pixel_y = int(round((y_max - y) * resolution_multiplier))
            polygon_coords.append((pixel_x, pixel_y))

        if len(polygon_coords) >= 3:
            draw.polygon(polygon_coords, fill=0)

    mask = np.array(mask_image, dtype=np.uint8) / 255

    return mask

def process_mask(file_path, mask_dir, lat0, lon0, h0, polygons):
    print(f"正在处理文件 {os.path.splitext(os.path.basename(file_path))[0]}")
    frames_info = load_frames_info(file_path)
    filename = frames_info["filename"]
    mask = create_mask(frames_info, lat0, lon0, h0, polygons, X_RANGE, Y_RANGE)

    mask_image = Image.fromarray((mask * 255).astype(np.uint8))
    mask_image.save(os.path.join(mask_dir, f"{filename}.png"))
    
def main():
    global X_RANGE, Y_RANGE

    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--X_RANGE", type=str, required=True)
    parser.add_argument("--Y_RANGE", type=str, required=True)
    args = parser.parse_args()
    
    X_RANGE = eval(args.X_RANGE)
    Y_RANGE = eval(args.Y_RANGE)

    folder_path = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR"
    label_file = os.path.join(output_folder, "label.json")
    meta_file = os.path.join(output_folder, "meta.json")
    mask_dir = os.path.join(output_folder, "label")
    
    os.makedirs(mask_dir, exist_ok=True)

    polygons = load_global_label(label_file)
    with open(meta_file, 'r', encoding='utf-8') as f:
        meta = json.load(f)
    polygons = convert_pixels_to_global(polygons, meta)

    lat0, lon0, h0 = meta['lat0'], meta['lon0'], meta['h0']

    thread = False
    if thread == True:
        with ThreadPoolExecutor() as executor:
            for root, _, files in os.walk(folder_path):
                files = sorted([f for f in files if f.endswith(".mat")], key=natural_key)
                for f in files:
                    file_path = os.path.join(root, f)
                    executor.submit(process_mask, file_path, mask_dir, lat0, lon0, h0, polygons)
    else:
        for root, _, files in os.walk(folder_path):
            files = sorted([f for f in files if f.endswith(".mat")], key=natural_key)
            for f in files:
                file_path = os.path.join(root, f)
                process_mask(file_path, mask_dir, lat0, lon0, h0, polygons)

if __name__ == "__main__":
    main()
